import torch
import torch.onnx
from model import BotRGCN

model = BotRGCN()  # 创建模型
model.load_state_dict(torch.load('train_module.pth'))
model.eval()  # 设置模型为评估模式

# 创建示例输入
des = torch.randn(229580, 768)  # 示例描述输入
tweet = torch.randn(229580, 768)  # 示例推文输入
num_prop = torch.randn(229580, 5)  # 示例数值特征输入（5个特征）
cat_prop = torch.randn(229580, 3)  # 示例类别特征输入（3个特征）

num_nodes = 229580  # 确保节点数量正确
num_edges = 10  # 边的数量

# 生成边索引，确保在有效范围内
edge_index = torch.randint(0, num_nodes, (2, num_edges)).long()
edge_type = torch.randint(0, 2, (num_edges,)).long()  # 转换边类型为 Long

# 导出为 ONNX
torch.onnx.export(model,
                  (des, tweet, num_prop, cat_prop, edge_index, edge_type),
                  'model.onnx',
                  input_names=['des', 'tweet', 'num_prop', 'cat_prop', 'edge_index', 'edge_type'],
                  output_names=['output'],
                  dynamic_axes={
                      'des': {0: 'batch_size'},
                      'tweet': {0: 'batch_size'},
                      'num_prop': {0: 'batch_size'},
                      'cat_prop': {0: 'batch_size'},
                      'edge_index': {0: 'num_edges'},
                      'edge_type': {0: 'num_edges'},
                      'output': {0: 'batch_size'}
                  })
